#include <atomic>
#include <curl/curl.h>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <vector>

constexpr int MAX_WORKERS = 8;
constexpr int PROGRESS_BAR_WIDTH = 50;

using std::cout;
using std::endl;
using std::string;

// 文件操作互斥锁
std::mutex file_mutex;

// 下载任务类
class DownloadTask {
public:
    std::string url;
    std::string filename;
    long start;
    long end;
    int worker_id;
    std::atomic<long> bytes_downloaded { 0 }; // 使用原子整数记录已下载字节数
    CURL* curl = nullptr;
    std::ofstream* file = nullptr; // 文件流指针

    DownloadTask(const DownloadTask&& dt)
    {
        url = dt.url;
        filename = dt.filename;
        start = dt.start;
        end = dt.end;
        worker_id = dt.worker_id;
    }

    DownloadTask(const std::string& url, const std::string& filename, long start,
        long end, int id)
        : url(url)
        , filename(filename)
        , start(start)
        , end(end)
        , worker_id(id)
    {
    }

    // 写入回调函数
    static size_t writeCallback(void* ptr, size_t size, size_t nmemb,
        void* userdata)
    {
        DownloadTask* task = static_cast<DownloadTask*>(userdata);
        size_t total_bytes = size * nmemb;

        // 加锁确保线程安全写入
        std::lock_guard<std::mutex> lock(file_mutex);

        if (task->file && task->file->is_open()) {
            // 计算写入位置
            long position = task->start + task->bytes_downloaded.load();

            // 移动到正确位置写入
            task->file->seekp(position);
            task->file->write(static_cast<char*>(ptr), total_bytes);

            // 更新已下载字节数
            task->bytes_downloaded += total_bytes;

            return total_bytes;
        }
        return 0;
    }

    // 进度回调函数
    static int progressCallback(void* clientp, curl_off_t dltotal,
        curl_off_t dlnow, curl_off_t ultotal,
        curl_off_t ulnow)
    {
        return 0; // 不再需要，使用writeCallback更精确
    }

    // 执行下载任务
    void execute()
    {
        // 打开文件并保持
        {
            std::lock_guard<std::mutex> lock(file_mutex);
            file = new std::ofstream(filename, std::ios::binary | std::ios::out);
            if (!file->is_open()) {
                std::cerr << "Worker " << worker_id << " failed to open file" << std::endl;
                return;
            }
        }

        curl = curl_easy_init();
        if (curl) {
            char range[100];
            snprintf(range, sizeof(range), "%ld-%ld", start, end);

            curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
            curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, &writeCallback);
            curl_easy_setopt(curl, CURLOPT_WRITEDATA, this);
            curl_easy_setopt(curl, CURLOPT_RANGE, range);
            curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
            curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 1L); // 禁用进度回调

            // 安全设置 - 仅用于测试环境
            curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, false);
            curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0);

            CURLcode res = curl_easy_perform(curl);
            if (res != CURLE_OK) {
                std::cerr << "Worker " << worker_id
                          << " failed: " << curl_easy_strerror(res) << std::endl;
            }

            // 清理资源
            curl_easy_cleanup(curl);
            {
                std::lock_guard<std::mutex> lock(file_mutex);
                if (file && file->is_open()) {
                    file->close();
                }
                delete file;
                file = nullptr;
            }
        }
    }

    // 获取下载进度百分比
    double getProgress() const
    {
        if (end <= start)
            return 1.0;
        return static_cast<double>(bytes_downloaded) / (end - start + 1);
    }
};

// 下载管理器类
class DownloadManager {
private:
    std::vector<DownloadTask> tasks;
    long total_size;
    std::atomic<long> downloaded;
    std::mutex output_mutex;
    bool finished = false;

public:
    DownloadManager(long total_size)
        : total_size(total_size)
        , downloaded(0)
    {
    }

    // 添加下载任务
    void addTask(const std::string& url, const std::string& filename, long start,
        long end, int id)
    {
        tasks.emplace_back(url, filename, start, end, id);
    }

    // 执行所有下载任务
    void executeAll()
    {
        std::vector<std::thread> threads;
        for (auto& task : tasks) {
            threads.emplace_back([&task]() { task.execute(); });
        }

        // 进度显示线程
        std::thread progress_thread([this]() {
            while (!finished) {
                updateDownloaded();
                showProgress();
                std::this_thread::sleep_for(std::chrono::seconds(1));
            }
        });

        // 等待所有下载线程完成
        for (auto& thread : threads) {
            thread.join();
        }

        finished = true;
        progress_thread.join();

        // 显示最终进度
        updateDownloaded();
        showProgress();
    }

    // 更新已下载量
    void updateDownloaded()
    {
        long total = 0;
        for (const auto& task : tasks) {
            total += task.bytes_downloaded;
        }
        downloaded = total;
    }

    // 显示下载进度
    void showProgress()
    {
        std::lock_guard<std::mutex> lock(output_mutex);
        // system("clear"); // 清屏
        cout << "Total size: " << formatSize(total_size) << endl;

        for (const auto& task : tasks) {
            int progress = static_cast<int>(task.getProgress() * 100);
            int bar_width = static_cast<int>(task.getProgress() * PROGRESS_BAR_WIDTH);

            std::cout << "Worker " << task.worker_id << " [" << std::setw(3)
                      << progress << "%] [";
            for (int j = 0; j < bar_width; j++)
                std::cout << "=";
            for (int j = bar_width; j < PROGRESS_BAR_WIDTH; j++)
                std::cout << " ";
            std::cout << "] "
                      << formatSize(task.bytes_downloaded)
                      << " / " << formatSize(task.end - task.start + 1) << std::endl;
        }

        double overall_progress = static_cast<double>(downloaded) / total_size;
        int overall_bar_width = static_cast<int>(overall_progress * PROGRESS_BAR_WIDTH);
        int overall_percent = static_cast<int>(overall_progress * 100);

        std::cout << "Overall [" << std::setw(3) << overall_percent << "%] [";
        for (int j = 0; j < overall_bar_width; j++)
            std::cout << "=";
        for (int j = overall_bar_width; j < PROGRESS_BAR_WIDTH; j++)
            std::cout << " ";
        std::cout << "] " << formatSize(downloaded) << " / "
                  << formatSize(total_size) << std::endl;
    }

    // 格式化文件大小显示
    static std::string formatSize(long bytes)
    {
        const char* units[] = { "B", "KB", "MB", "GB" };
        int unit = 0;
        double size = bytes;
        while (size >= 1024 && unit < 3) {
            size /= 1024;
            unit++;
        }
        std::ostringstream oss;
        oss << std::fixed << std::setprecision(2) << size << " " << units[unit];
        return oss.str();
    }
};

void printSupportedProtocols()
{
    curl_version_info_data* ver = curl_version_info(CURLVERSION_NOW);
    std::cout << "Supported protocols: ";
    for (int i = 0; ver->protocols[i]; i++) {
        std::cout << ver->protocols[i] << " ";
    }
    std::cout << std::endl;
}

// 获取文件大小
long getFileSize(const std::string& url)
{
    printSupportedProtocols();
    CURL* curl = curl_easy_init();
    long file_size = 0;
    if (curl) {
        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl, CURLOPT_NOBODY, 1L);
        curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
        curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, false);
        curl_easy_setopt(curl, CURLOPT_SSL_VERIFYHOST, 0);

        CURLcode res = curl_easy_perform(curl);
        if (res == CURLE_OK) {
            curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD_T, &file_size);
        } else {
            std::cerr << "Failed to get file size: " << curl_easy_strerror(res)
                      << std::endl;
        }
        // curl_easy_cleanup(curl); // 修复：添加清理
    }
    return file_size;
}

// 获取已存在文件大小
long getExistingFileSize(const std::string& filename)
{
    std::ifstream file(filename, std::ios::binary | std::ios::ate);
    if (file) {
        return file.tellg();
    }
    return 0;
}

int main(int argc, char* argv[])
{
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0]
                  << " <URL> [workers=" << MAX_WORKERS << "]"
                  << std::endl;
        return 1;
    }

    std::string url = argv[1];
    auto first = url.find("github.com");

    std::string ghproxy = "https://gh-proxy.com/";
    if (first != std::string::npos) {
        url = ghproxy + url.substr(first);
    }

    int pos = url.find_last_of('/') + 1;
    std::string filename = url.substr(pos);
    // 确保文件名有效
    if (filename.empty() || filename == "/" || filename == ".") {
        filename = "downloaded_file";
    }

    int workers = (argc > 2) ? std::stoi(argv[2]) : MAX_WORKERS;
    if (workers <= 0 || workers > MAX_WORKERS) {
        workers = MAX_WORKERS;
    }

    // 初始化libcurl
    curl_global_init(CURL_GLOBAL_ALL);
    // 获取文件大小
    long file_size = getFileSize(url);
    if (file_size <= 0) {
        std::cerr << "Could not determine file size or file is empty" << std::endl;
        curl_global_cleanup();
        return 1;
    }

    // 检查文件是否已存在，获取已下载大小
    long existing_size = getExistingFileSize(filename);

    // 如果已有部分下载，启用断点续传
    if (existing_size > 0 && existing_size < file_size) {
        std::cout << "Resuming download from "
                  << DownloadManager::formatSize(existing_size) << std::endl;
    } else if (existing_size >= file_size) {
        std::cout << "File already downloaded" << std::endl;
        curl_global_cleanup();
        return 0;
    } else {
        existing_size = 0;
        // 创建空文件
        {
            std::ofstream file(filename, std::ios::binary);
            if (!file) {
                std::cerr << "Failed to create file: " << filename << std::endl;
                curl_global_cleanup();
                return 1;
            }
        }
    }

    // 创建下载管理器
    DownloadManager manager(file_size);

    // 计算每个线程的下载范围
    long remaining_size = file_size - existing_size;
    long chunk_size = remaining_size / workers;

    // 重新分配任务，确保不重叠且覆盖整个文件
    for (int i = 0; i < workers; i++) {
        long start = existing_size + i * chunk_size;
        long end = (i == workers - 1) ? file_size - 1 : existing_size + (i + 1) * chunk_size - 1;

        // 只添加有数据需要下载的任务
        if (start < end) {
            manager.addTask(url, filename, start, end, i);
        }
    }

    // 执行下载
    manager.executeAll();

    // 清理
    curl_global_cleanup();

    std::cout << "Download complete!" << std::endl;

    return 0;
}
